from torchvision.datasets import CIFAR10, CIFAR100
from torchvision.transforms import transforms
from torch.utils.data import DataLoader, Dataset
import torchvision

cifar10_mean = (0.4914, 0.4822, 0.4465)
cifar10_std = (0.2471, 0.2435, 0.2616)

cifar100_mean = (0.5070751592371323,
                 0.48654887331495095, 0.4409178433670343)
cifar100_std = (0.2673342858792401, 0.2564384629170883,
                0.27615047132568404)

cifar10_train_transform = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.2, 1.)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([
        transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened
    ], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize(cifar10_mean, cifar10_std),
    # transforms.RandomErasing(p=0.25)
])
cifar10_test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(cifar10_mean, cifar10_std)
])

mae_cifar10_train_trainsform = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.2, 1.)),
    transforms.ToTensor(),
    transforms.Normalize(cifar10_mean, cifar10_std)
])

mae_cifar100_train_trainsform = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.2, 1.)),
    transforms.ToTensor(),
    transforms.Normalize(cifar100_mean, cifar100_std)
])

cifar100_train_transform = transforms.Compose([
    transforms.RandomResizedCrop(32, scale=(0.2, 1.)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([
        transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened
    ], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize(cifar100_mean, cifar100_std),
    # transforms.RandomErasing(p=0.25)
])
cifar100_test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(cifar100_mean, cifar100_std)
])


def normalize(dataset):
    return {
        "imagenet": [[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]],
        "cifar10": [[0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]],
        "cifar100": [[0.5070751592371323, 0.48654887331495095, 0.4409178433670343], [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]],
        "tiny_imagenet": [[0.4802, 0.4481, 0.3975], [0.2770, 0.2691, 0.2821]]
    }[dataset]

# class TransformTwice:
#     def __init__(self, transform):
#         self.transform = transform

#     def __call__(self, inp):
#         out1 = self.transform(inp)
#         out2 = self.transform(inp)
#         return out1, out2
# class TransformTwice:
#     def __init__(self, transform):
#         self.transform = transform

#     def __call__(self, inp):
#         out1 = self.transform(inp)
#         out2 = self.transform(inp)
#         return out1, out2


def get_image_size(args):
    if args.dataset == "imagenet" or args.dataset == "tiny_imagenet":
        image_size = 224
    elif args.dataset == "cifar10" or args.dataset == "cifar100":
        image_size = 32
    else:
        raise NotImplementedError
    return image_size


class TransformTwice:
    def __init__(self,args):
        image_size=get_image_size(args)
        p_blur = 0.5 if image_size > 32 else 0
        self.transform = transforms.Compose([
            transforms.RandomResizedCrop(image_size, scale=(0.2, 1.0)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply(
                [transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            # transforms.RandomApply([transforms.GaussianBlur([.1, 2.])], p=p_blur),
            transforms.ToTensor(),
            transforms.Normalize(*normalize(args.dataset))
        ])

    def __call__(self, x):
        x1 = self.transform(x)
        x2 = self.transform(x)
        return x1, x2


def test_transform(args):
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(*normalize(args.dataset))
    ])


def test_crop_transform(args):
    return transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(*normalize(args.dataset))
    ])


def mae_transform(args):
    return transforms.Compose([
        transforms.RandomResizedCrop(32, scale=(0.2, 1.)),
        transforms.ToTensor(),
        transforms.Normalize(*normalize(args.dataset))
    ])


def train_transform(args):
    image_size=get_image_size(args)

    return transforms.Compose([
        transforms.RandomResizedCrop(image_size, scale=(0.2, 1.)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomApply([
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)  # not strengthened
        ], p=0.8),
        transforms.RandomGrayscale(p=0.2),
        transforms.ToTensor(),
        transforms.Normalize(*normalize(args.dataset)),
    ])

# class TransformTwice:
#     def __init__(self, image_size,args):
#         image_size = 224 if image_size is None else image_size
#         p_blur = 0.5 if image_size > 32 else 0
#         self.transform = transforms.Compose([
#             transforms.RandomResizedCrop(image_size, scale=(0.2, 1.0)),
#             transforms.RandomHorizontalFlip(),
#             transforms.RandomApply(
#                 [transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
#             transforms.RandomGrayscale(p=0.2),
#             transforms.RandomApply(
#                 [transforms.GaussianBlur([.1, 2.])], p=p_blur),
#             transforms.ToTensor(),
#             transforms.Normalize(*normalize(args.dataset))
#         ])

#     def __call__(self, x):
#         x1 = self.transform(x)
#         x2 = self.transform(x)
#         return x1, x2
